import os  # 用于文件路径和操作系统相关操作
import time  # 用于计时
import random  # 用于随机数生成
import json  # 用于JSON文件读写
from tqdm import tqdm  # 用于显示进度条（当前代码未直接使用，可能预留）

import torch  # PyTorch深度学习框架
import torch.nn as nn  # 神经网络模块
from torch.nn.utils import clip_grad_norm_  # 用于梯度裁剪，防止梯度爆炸
from torch.utils.data import DataLoader  # 数据加载器，用于批量加载数据
from tensorboardX import SummaryWriter  # 用于TensorBoard可视化
from torch.optim.lr_scheduler import StepLR, MultiStepLR  # 学习率调度器

import numpy as np  # 用于数值计算
from configs.opts import parser  # 从配置文件导入参数解析器
from model.main_model import supv_main_model as main_model  # 导入主模型（监督学习版本）
from utils import AverageMeter, Prepare_logger, get_and_save_args  # 导入工具类：平均计量器、日志准备、参数处理
from utils.Recorder import Recorder  # 用于记录实验参数和结果
from dataset.AVE_dataset import AVEDataset  # 导入AVE数据集类
import torch.nn.functional as F  # 导入PyTorch的函数式接口


# =================================  随机种子配置  ============================
# 设置随机种子，保证实验的可重复性
SEED = 43
random.seed(SEED)  # Python随机数种子
np.random.seed(seed=SEED)  # NumPy随机数种子
torch.manual_seed(seed=SEED)  # PyTorch CPU随机数种子
torch.cuda.manual_seed(seed=SEED)  # PyTorch GPU随机数种子
torch.backends.cudnn.deterministic = True  # 确保CUDA卷积操作的确定性
torch.backends.cudnn.benchmark = False  # 关闭CUDA的基准测试模式（可能影响速度，但保证确定性）


# 加载模型配置文件
config_path = 'configs/main.json'  # 配置文件路径
with open(config_path) as fp:
    config = json.load(fp)  # 加载JSON格式的配置
print(config)  # 打印配置信息，确认加载正确
# =============================================================================


def AVPSLoss(av_simm, soft_label):
    """
    音频-视觉对相似度损失（用于全监督设置）
    对应论文中的公式(8, 9)
    参数:
        av_simm: 音频-视觉相似度矩阵，形状为[batch_size, 10]
        soft_label: 软标签，用于计算MSE损失
    返回:
        计算得到的损失值
    """
    relu_av_simm = F.relu(av_simm)  # 对相似度取ReLU，过滤负值
    sum_av_simm = torch.sum(relu_av_simm, dim=-1, keepdim=True)  # 沿最后一维求和，用于归一化
    avg_av_simm = relu_av_simm / (sum_av_simm + 1e-8)  # 归一化得到平均相似度（加小值防止除零）
    loss = nn.MSELoss()(avg_av_simm, soft_label)  # 计算与软标签的MSE损失
    return loss


def main():
    # 全局变量定义
    global args, logger, writer, dataset_configs
    # 用于记录最佳准确率及对应epoch
    global best_accuracy, best_accuracy_epoch
    best_accuracy, best_accuracy_epoch = 0, 0  # 初始化最佳准确率为0，对应epoch为0

    # 加载并处理数据集配置参数
    dataset_configs = get_and_save_args(parser)  # 从解析器获取并保存数据集配置
    parser.set_defaults(** dataset_configs)  # 将数据集配置设置为解析器默认值
    args = parser.parse_args()  # 解析命令行参数

    # 设置使用的GPU设备
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"  # 按PCI总线ID排序GPU
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu  # 指定使用的GPU设备

    '''创建快照目录，用于保存代码副本和模型文件'''
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)  # 若目录不存在则创建

    # 若指定了恢复路径，将快照目录设为恢复路径的父目录
    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    # 准备日志记录器（根据是否为评估模式设置不同日志）
    logger = Prepare_logger(args, eval=args.evaluate)

    # 若不是评估模式，记录创建目录和运行参数；否则记录评估日志路径
    if not args.evaluate:
        logger.info(f'\nCreating folder: {args.snapshot_pref}')
        logger.info('\nRuntime args\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    else:
        logger.info(f'\nLog file will be save in a {args.snapshot_pref}/Eval.log.')

    '''数据集加载'''
    # 训练集数据加载器
    train_dataloader = DataLoader(
        AVEDataset('./data/', split='train'),  # 加载训练集
        batch_size=args.batch_size,  # 批次大小
        shuffle=True,  # 训练时打乱数据
        num_workers=8,  # 数据加载的进程数
        pin_memory=True  # 锁存内存，加速GPU数据传输
    )

    # 测试集数据加载器
    test_dataloader = DataLoader(
        AVEDataset('./data/', split='test'),  # 加载测试集
        batch_size=args.test_batch_size,  # 测试批次大小
        shuffle=False,  # 测试时不打乱数据
        num_workers=8,  # 数据加载的进程数
        pin_memory=True  # 锁存内存
    )

    '''模型设置'''
    mainModel = main_model(config['model'])  # 初始化主模型（使用配置文件中的模型参数）
    mainModel = nn.DataParallel(mainModel).cuda()  # 数据并行（多GPU训练），并转移到GPU
    learned_parameters = mainModel.parameters()  # 获取模型可学习参数
    optimizer = torch.optim.Adam(learned_parameters, lr=args.lr)  # 定义Adam优化器
    # 学习率调度器：多步调度，在指定epoch处将学习率乘以gamma
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.5)
    # 定义损失函数：BCEWithLogitsLoss（带sigmoid的二分类损失）和CrossEntropyLoss（多分类损失）
    criterion = nn.BCEWithLogitsLoss().cuda()
    criterion_event = nn.CrossEntropyLoss().cuda()

    '''从检查点恢复模型（若指定）'''
    if os.path.isfile(args.resume):
        logger.info(f"\nLoading Checkpoint: {args.resume}\n")
        mainModel.load_state_dict(torch.load(args.resume))  # 加载模型权重
    elif args.resume != "" and (not os.path.isfile(args.resume)):
        raise FileNotFoundError  # 若指定了恢复路径但文件不存在，抛出错误

    '''仅评估模式（若指定）'''
    if args.evaluate:
        logger.info(f"\nStart Evaluation..")
        validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch=0, eval_only=True)
        return  # 评估完成后退出

    '''TensorBoard可视化和代码备份'''
    writer = SummaryWriter(args.snapshot_pref)  # 初始化TensorBoard写入器
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")  # 初始化记录器，用于备份代码
    recorder.writeopt(args)  # 记录实验参数

    '''训练和测试循环'''
    for epoch in range(args.n_epoch):  # 遍历每个epoch
        # 训练一个epoch，返回平均损失
        loss = train_epoch(mainModel, train_dataloader, criterion, criterion_event, optimizer, epoch)

        # 按指定频率或最后一个epoch进行验证
        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            acc = validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch)
            # 若当前准确率高于最佳准确率，更新最佳记录并保存模型
            if acc > best_accuracy:
                best_accuracy = acc
                best_accuracy_epoch = epoch
                save_checkpoint(
                    mainModel.state_dict(),
                    top1=best_accuracy,
                    task='Supervised',
                    epoch=epoch + 1,
                )
            print("-----------------------------")
            print("best acc and epoch:", best_accuracy, best_accuracy_epoch)
            print("-----------------------------")
        scheduler.step()  # 每个epoch更新学习率


def train_epoch(model, train_dataloader, criterion, criterion_event, optimizer, epoch):
    """
    训练一个epoch的函数
    参数:
        model: 待训练的模型
        train_dataloader: 训练数据加载器
        criterion: 二分类损失函数（BCE）
        criterion_event: 多分类损失函数（CrossEntropy）
        optimizer: 优化器
        epoch: 当前epoch编号
    返回:
        该epoch的平均损失
    """
    # 初始化计量器：批次时间、数据加载时间、损失、训练准确率
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    train_acc = AverageMeter()
    end_time = time.time()  # 记录起始时间

    model.train()  # 设置模型为训练模式（启用 dropout、BN等）
    model.double()  # 将模型精度设为double（因输入特征为double类型）
    optimizer.zero_grad()  # 初始化梯度为0

    for n_iter, batch_data in enumerate(train_dataloader):  # 遍历每个批次
        data_time.update(time.time() - end_time)  # 更新数据加载时间

        '''输入数据到模型'''
        visual_feature, audio_feature, labels = batch_data  # 解包批次数据（视觉特征、音频特征、标签）
        labels = labels.double().cuda()  # 将标签转为double类型并转移到GPU

        # 模型前向传播：输出事件存在分数、事件类别分数、音视频门控、音视频分数
        is_event_scores, event_scores, audio_visual_gate, av_score = model(visual_feature, audio_feature)
        # 调整输出维度以匹配标签（转置并压缩维度）
        is_event_scores = is_event_scores.transpose(1, 0).squeeze().contiguous()
        audio_visual_gate = audio_visual_gate.transpose(1, 0).squeeze().contiguous()

        # 处理标签：提取前景标签（排除背景），获取二分类标签和事件类别标签
        labels_foreground = labels[:, :, :-1]  # [32, 10, 28]（假设28为背景类别）
        labels_BCE, labels_evn = labels_foreground.max(-1)  # 二分类标签（是否有事件）和事件类别标签
        labels_event, _ = labels_evn.max(-1)  # 最终事件类别标签

        # 计算各部分损失
        loss_is_event = criterion(is_event_scores, labels_BCE.cuda())  # 事件存在性损失（二分类）
        label_is_gate = criterion(audio_visual_gate, labels_BCE.cuda())  # 音视频门控损失（二分类）
        loss_cas = criterion_event(av_score, labels_event.cuda())  # 音视频分数的事件分类损失
        loss_event_class = criterion_event(event_scores, labels_event.cuda())  # 事件类别分类损失

        # 总损失：各部分损失相加
        loss = loss_is_event + label_is_gate + loss_event_class + loss_cas
        loss.backward()  # 反向传播计算梯度

        '''计算准确率'''
        acc = compute_accuracy_supervised(is_event_scores, event_scores, labels)
        train_acc.update(acc.item(), visual_feature.size(0) * 10)  # 更新准确率计量器（batch_size*10为总样本数）

        '''梯度裁剪（防止梯度爆炸）'''
        if args.clip_gradient is not None:
            total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)  # 裁剪梯度

        '''更新参数'''
        optimizer.step()  # 优化器更新参数
        optimizer.zero_grad()  # 清空梯度，准备下一批次

        # 更新损失计量器
        losses.update(loss.item(), visual_feature.size(0) * 10)
        batch_time.update(time.time() - end_time)  # 更新批次时间
        end_time = time.time()  # 重置结束时间

        '''将迭代损失写入TensorBoard'''
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        '''在终端打印日志（按指定频率）'''
        if n_iter % args.print_freq == 0:
            logger.info(
                f'Train Epoch: [{epoch}][{n_iter}/{len(train_dataloader)}]\t'
                f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                f'Prec@1 {train_acc.val:.3f} ({train_acc.avg: .3f})'
            )

        '''将epoch平均损失写入TensorBoard'''
        writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    return losses.avg  # 返回该epoch的平均损失


@torch.no_grad()  # 禁用梯度计算（节省内存，加速推理）
def validate_epoch(model, test_dataloader, criterion, criterion_event, epoch, eval_only=False):
    """
    验证一个epoch的函数（与训练类似，但无参数更新）
    参数:
        model: 待验证的模型
        test_dataloader: 测试数据加载器
        criterion: 二分类损失函数
        criterion_event: 多分类损失函数
        epoch: 当前epoch编号
        eval_only: 是否仅评估模式
    返回:
        验证集的平均准确率
    """
    # 初始化计量器
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    accuracy = AverageMeter()
    end_time = time.time()

    model.eval()  # 设置模型为评估模式（关闭dropout、BN固定）
    model.double()  # 同训练，使用double精度

    for n_iter, batch_data in enumerate(test_dataloader):  # 遍历测试集批次
        data_time.update(time.time() - end_time)

        '''输入数据到模型'''
        visual_feature, audio_feature, labels = batch_data
        labels = labels.double().cuda()
        bs = visual_feature.size(0)  # 获取批次大小

        # 模型前向传播（不计算梯度）
        is_event_scores, event_scores, audio_visual_gate, _ = model(visual_feature, audio_feature)
        is_event_scores = is_event_scores.transpose(1, 0).squeeze()
        audio_visual_gate = audio_visual_gate.transpose(1, 0).squeeze()

        # 处理标签（同训练）
        labels_foreground = labels[:, :, :-1]
        labels_BCE, labels_evn = labels_foreground.max(-1)
        labels_event, _ = labels_evn.max(-1)

        # 计算验证损失
        loss_is_event = criterion(is_event_scores, labels_BCE.cuda())
        loss_is_gate = criterion(audio_visual_gate, labels_BCE.cuda())
        loss_event_class = criterion_event(event_scores, labels_event.cuda())
        loss = loss_is_event + loss_event_class + loss_is_gate

        # 计算准确率
        acc = compute_accuracy_supervised(is_event_scores, event_scores, labels)
        accuracy.update(acc.item(), bs * 10)  # 更新准确率计量器

        # 更新时间和损失计量器
        batch_time.update(time.time() - end_time)
        end_time = time.time()
        losses.update(loss.item(), bs * 10)

        '''打印验证日志（按指定频率）'''
        if n_iter % args.print_freq == 0:
            logger.info(
                f'Test Epoch [{epoch}][{n_iter}/{len(test_dataloader)}]\t'
                f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                f'Prec@1 {accuracy.val:.3f} ({accuracy.avg:.3f})'
            )

    '''将验证损失和准确率写入TensorBoard（非仅评估模式）'''
    if not eval_only:
        writer.add_scalar('Val_epoch_data/epoch_loss', losses.avg, epoch)
        writer.add_scalar('Val_epoch/Accuracy', accuracy.avg, epoch)

    # 记录验证集整体准确率
    logger.info(
        f'**************************************************************************\t'
        f"\tEvaluation results (acc): {accuracy.avg:.4f}%."
    )
    return accuracy.avg  # 返回平均准确率


def compute_accuracy_supervised(is_event_scores, event_scores, labels):
    """
    计算监督学习模式下的准确率
    参数:
        is_event_scores: 事件存在性分数（是否有事件）
        event_scores: 事件类别分数
        labels: 真实标签
    返回:
        计算得到的准确率（百分比）
    """
    _, targets = labels.max(-1)  # 从标签中提取真实类别（取最大值索引）

    # 事件存在性预测：sigmoid后大于0.5视为有事件
    is_event_scores = is_event_scores.sigmoid()
    scores_pos_ind = is_event_scores > 0.5  # 有事件的掩码
    scores_mask = scores_pos_ind == 0  # 无事件（背景）的掩码

    # 事件类别预测：取分数最高的类别
    _, event_class = event_scores.max(-1)
    pred = scores_pos_ind.long()  # 将存在性预测转为长整型
    pred *= event_class[:, None]  # 结合存在性和类别（存在事件时为类别，否则为0）

    # 无事件时预测为背景（28表示背景类别）
    pred[scores_mask] = 28

    # 计算准确率：预测与真实标签一致的比例
    correct = pred.eq(targets)
    correct_num = correct.sum().double()
    acc = correct_num * (100. / correct.numel())  # 转为百分比

    return acc


def save_checkpoint(state_dict, top1, task, epoch):
    """
    保存模型检查点
    参数:
        state_dict: 模型状态字典
        top1: 最佳准确率
        task: 任务名称（此处为监督学习）
        epoch: 对应epoch
    """
    # 构建模型保存路径（包含epoch、准确率、任务名称）
    model_name = f'{args.snapshot_pref}/model_epoch_{epoch}_top1_{top1:.3f}_task_{task}_best_model.pth.tar'
    torch.save(state_dict, model_name)  # 保存模型权重


if __name__ == '__main__':
    main()  # 程序入口
